import numpy as np
import matplotlib.pyplot as plt
from disba import PhaseDispersion
from dispCal.disp import calDisp
import bayesbay as bb
from bayesbay.discretization import Voronoi1D
from bayesbay.likelihood import LogLikelihood, Target
import os
import glob
from functools import partial
from utils.h5data import read_h5_data

# 创建结果保存根目录
output_root_dir = "figures/6.traceSeq.figures/test.inv"
os.makedirs(output_root_dir, exist_ok=True)

VP_VS = 1.77
RHO_VP = [0.32, 0.77]

def read_dispersion_curves(filepath):
    """
    从文件中读取频散曲线数据
    """
    with open(filepath, 'r') as f:
        lines = f.readlines()
    
    curves = []
    current_curve = None
    
    for line in lines:
        line = line.strip()
        if line.startswith('#'):
            if 'Curve' in line:
                if current_curve is not None and len(current_curve)>=1:
                    curves.append(np.array(current_curve))
                current_curve = []
        else:
            if line and current_curve is not None:
                freq, vel = map(float, line.split())
                current_curve.append([freq, vel])
    
    # 添加最后一个curve
    if current_curve is not None and len(current_curve)>=1:
        curves.append(np.array(current_curve))
    
    return curves

def forward_sw(state, periods, vp_vs, rho_vp, wave="rayleigh", mode=0):
    """
    正演函数
    """
    voronoi = state["voronoi"]
    voronoi_sites = voronoi["discretization"]
    thickness = Voronoi1D.compute_cell_extents(voronoi_sites)
    vs = voronoi["vs"]
    if np.min(vs[1:]-vs[:-1])<-0.07:
        return np.zeros_like(periods)
    
    vp = vs * vp_vs
    rho = rho_vp[0] * vp + rho_vp[1]
    
    # 计算瑞雷波相速度
    # pd = PhaseDispersion(thickness, vp, vs, rho)
    # d_pred_rayleigh = pd(periods, mode=mode, wave="rayleigh").velocity

    d_pred_rayleigh = calDisp(thickness,vp,vs,rho,periods,mode=mode+1)

    if len(d_pred_rayleigh) < len(periods):
        re = np.zeros_like(len(periods)-len(d_pred_rayleigh))+max(vs)
        d_pred_rayleigh = np.concatenate([d_pred_rayleigh, re])

    return d_pred_rayleigh

def compute_dispersion_curves(thickness, vs, vp_vs=VP_VS, rho_vp=RHO_VP):
    """
    根据给定的模型参数计算频散曲线
    """
    # 确保厚度数组正确（最后一个为半空间）
    thickness = thickness.copy()
    if len(thickness) > 0:
        thickness[-1] = 0  # 最后一层为半空间
    
    vp = vs * vp_vs
    rho = rho_vp[0] * vp + rho_vp[1]
    
    # 计算瑞雷波相速度
    pd = PhaseDispersion(thickness, vp, vs, rho)
    # 根据观测数据的频率范围计算理论曲线
    modes = []
    f = np.linspace(20,1,100)
    periods = 1/f
    for n in range(10):
        try:
            # 计算n阶模式
            # r = pd(periods, mode=n, wave="rayleigh")
            # modes.append((r.period, r.velocity, f'R{n}'))

            c = calDisp(thickness,vp,vs,rho,periods,mode=n+1)
            modes.append((periods, c, f'R{n}'))
        except Exception as e:
            print(f"Error computing R0 mode: {e}")
    print(f"Computed {len(modes)} modes")
        
    print(modes[-1])    
    return modes

def initialize_vs(param, positions=None):
    """
    初始化Vs参数
    """
    # vmin, vmax = param.get_vmin_vmax(positions)
    vmin, vmax = param.get_vmin_vmax(positions)
    sorted_vals = np.sort(np.random.uniform(vmin, vmax, positions.size))
    return sorted_vals

def create_parameterization():
    """
    创建参数化模型
    """
    # vs = bb.prior.UniformPrior(name="vs", 
    #                            vmin=[0.37, 0.4, 0.6], 
    #                            vmax=[0.5, 0.6, 1.4], 
    #                            position=[0, 0.1, 0.2],
    #                            perturb_std=0.03)
    # vs.set_custom_initialize(initialize_vs)
    # vs = bb.prior.GaussianPrior(name="vs", 
    #                            mean=[0.37, 0.6, 1.2], 
    #                            std=[0.05, 0.08, 0.2], 
    #                            position=[0, 0.05, 0.3],
    #                            perturb_std=0.03)
    vs = bb.prior.GaussianPrior(name="vs", 
                               mean=[0.37, 0.6, 1.2], 
                               std=[0.05, 0.1, 0.2], 
                               position=[0, 0.05, 0.3],
                               perturb_std=0.03)
    
    
    vp_vs = VP_VS  # vp/vs比率
    rho_vp = RHO_VP  # rho = 0.32 * vp + 0.77
    Zmin, Zmax = 0, 0.6
    
    voronoi = Voronoi1D(
        name="voronoi", 
        vmin=Zmin,   #最小深度
        vmax=Zmax,   #最大深度
        perturb_std=0.01,
        n_dimensions=None, 
        n_dimensions_min=3,
        n_dimensions_max=10,
        parameters=[vs], 
        birth_from='neighbour'
    )
    
    parameterization = bb.parameterization.Parameterization(voronoi)
    return parameterization, vp_vs, rho_vp, Zmin, Zmax

def perform_inversion(curves_data, output_dir, filename, fv_file = None):
    """
    对给定的曲线数据执行反演
    """
    os.makedirs(output_dir, exist_ok=True)
    
    # 创建参数化模型
    parameterization, vp_vs, rho_vp, Zmin, Zmax = create_parameterization()
    print("参数化模型:", parameterization)
    print("Zmin, Zmax:", Zmin, Zmax)
    # 处理曲线数据
    n_curves = len(curves_data)
    if n_curves == 0:
        print(f"警告: 文件 {filename} 中没有曲线数据")
        return
    
    # 提取第一阶模式数据
    mode0_data = curves_data[0]
    freq0, d_r0 = mode0_data[-1::-1, 0], mode0_data[-1::-1, 1]/1000
    periods0 = 1.0 / freq0
    
    targets = []
    fwd_functions = []
    
    # 添加第一阶模式
    target_rayleigh0 = Target("R0", 
                                d_r0, 
                                std_min=0.005, 
                                std_max=0.02, 
                                std_perturb_std=0.005)
    f0 = partial(forward_sw, periods=periods0, vp_vs=vp_vs, rho_vp=rho_vp, wave="rayleigh", mode=0)
    targets.append(target_rayleigh0)
    fwd_functions.append(f0)
    print("添加基阶模式")
    print(f'fmax={freq0[0]}, fmin={freq0[-1]}, nf={len(freq0)}')

    if n_curves > 1:
        for i in range(1, n_curves):
            mode_data = curves_data[i]
            freqi, d_ri = mode_data[-1::-1, 0], mode_data[-1::-1, 1]/1000
            periodsi = 1.0 / freqi
            
            target_rayleighi = Target(f"R{i}", 
                                        d_ri, 
                                        std_min=0.01, 
                                        std_max=0.04, 
                                        std_perturb_std=0.01)
            f1 = partial(forward_sw, periods=periodsi, vp_vs=vp_vs, rho_vp=rho_vp, wave="rayleigh", mode=i)
            targets.append(target_rayleighi)
            fwd_functions.append(f1)

            print(f"添加第{i}阶模式")
            print(f'fmax={freqi[0]}, fmin={freqi[-1]}, nf={len(freqi)}')
        
    
    # 创建似然函数
    log_likelihood = LogLikelihood(targets=targets, fwd_functions=fwd_functions)
    
    # 贝叶斯反演
    inversion = bb.BayesianInversion(
        parameterization=parameterization, 
        log_likelihood=log_likelihood,
        n_chains=16
    )
    
    print(f"开始反演文件 {filename}...")
    print(f"观测数据参数:")
    print(f"  R0: {periods0.shape}, {d_r0.shape}")
    
    inversion.run(
        sampler=None, 
        n_iterations=100_000, 
        burnin_iterations=20_000, 
        save_every=100,
        verbose=False,
        print_every=10000,
    )
    
    # 输出统计信息
    for chain in inversion.chains:
        chain.print_statistics()
    
    # 获取结果
    results = inversion.get_results(concatenate_chains=True)
    
    # 处理结果
    sampled_voronoi_nuclei = results['voronoi.discretization']
    sampled_thickness = [Voronoi1D.compute_cell_extents(n) for n in sampled_voronoi_nuclei]
    sampled_vs = results['voronoi.vs']
    interp_depths = np.linspace(Zmin, Zmax, 100)
    statistics_vs = Voronoi1D.get_tessellation_statistics(
        sampled_thickness, sampled_vs, interp_depths, input_type='extents'
    )
    
    sampled_thickness_valid = []
    for i in sampled_thickness:
        ti = list(i[:-1])+[0.05]
        sampled_thickness_valid.append(np.array(ti))
    sampled_thickness = sampled_thickness_valid
    
    # 绘制结果
    plt.close('all')
    fig, axes = plt.subplots(1, 2, figsize=(12, 6))

    # 读取FV数据
    if fv_file is not None:
        FV, f_vals, vs_vals = read_h5_data(fv_file,['FV','f','vs'])
        # 在背景上绘制FV能量分布
        X,Y = np.meshgrid(f_vals, vs_vals/1000)
        axes[0].pcolormesh(X,Y, FV.T, cmap='Greys')
        axes[0].set_ylim([vs_vals.min()/1000, vs_vals.max()/1000])
        # axes[0].set_xlim([f_vals.min(), f_vals.max()])
        axes[0].set_xlim([1,6])
        

    
    # 绘制频散曲线对比
    colors = ['red','blue','green']
    for i in range(0, n_curves):
        mode_data = curves_data[i]
        freqi, d_ri = mode_data[-1::-1, 0], mode_data[-1::-1, 1]/1000
        periodsi = 1.0 / freqi

        axes[0].plot(1/periodsi, d_ri, 'o',color=colors[i], label=f'R{i}', markerfacecolor='none')

        d_pred_rayleigh_all = np.array(results[f'R{i}.dpred'])
        
        # r_mean = np.mean(d_pred_rayleigh_all, axis=0)
        # axes[0].plot(1/periodsi, r_mean, linewidth=1, 
        #                         color=colors[i], label=f'R{i} pred')
        
        percentiles_rayleigh = np.percentile(d_pred_rayleigh_all, (10, 90), axis=0)
        axes[0].fill_between(1/periodsi, percentiles_rayleigh[0], percentiles_rayleigh[1], 
                                color=colors[i], alpha=0.5, label=None)
    # 计算理论频散曲线
    thickness = np.diff(interp_depths)
    thickness = np.append(thickness, 0)  # 最后一层为半空间
    mean_vs = statistics_vs['median']
    theoretical_modes = compute_dispersion_curves(thickness, mean_vs)
    # 绘制理论频散曲线
    for i, (period, velocity, mode_name) in enumerate(theoretical_modes):
        if i < len(colors):
            axes[0].plot(1/period, velocity, color=colors[i], linestyle='--', 
                     label=f'R{i} median', linewidth=1)
        else:
            axes[0].plot(1/period, velocity, color='orange',   linestyle='-', 
                     label=None, linewidth=0.7)

    
    axes[0].set_xlabel('F(Hz)')
    axes[0].set_ylabel('Phase velocity [km/s]')
    axes[0].legend()
    axes[0].grid(True)
    axes[0].set_title('Dispersion Curves')
    
    # 绘制Vs结构
    ax1 = axes[1]
    Voronoi1D.plot_tessellation_density(sampled_thickness, sampled_vs, 
                                        input_type='extents', ax=ax1, cmap='gray_r')

    ax1.plot(statistics_vs['mean'], interp_depths, 'r-', lw=2, label='Mean')
    ax1.plot(statistics_vs['median'], interp_depths, 'b--', lw=2, label='Median')
    ax1.set_xlabel('Vs [km/s]')
    ax1.set_ylabel('Depth [km]')
    ax1.legend()
    ax1.grid(True)
    ax1.set_title('Inverted Vs Profile')
    ax1.set_ylim(Zmax,0)
    ax1.set_xlim(0, 1.5)
    
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f"{filename}_inversion_results.png"), dpi=300, bbox_inches='tight')
    plt.close()

    # assert False
    
    # 绘制链结果
    def get_subplot_layout(n_subplots):
        rows = int(np.sqrt(n_subplots))
        cols = int(np.ceil(n_subplots / rows))
        return rows, cols
    
    rows, cols = get_subplot_layout(len(inversion.chains))
    fig, axes = plt.subplots(rows, cols, figsize=(15, 15))
    axes = np.ravel(axes)
    
    for ipanel, (ax, chain) in enumerate(zip(axes, inversion.chains)):
        saved_states = chain.saved_states
        saved_nuclei = saved_states["voronoi.discretization"]
        saved_vs = saved_states['voronoi.vs']
        
        Voronoi1D.plot_tessellations(saved_nuclei,
                                     saved_vs, 
                                     ax=ax, 
                                     linewidth=0.1, 
                                     color="k", 
                                     bounds=(0, Zmax))
        Voronoi1D.plot_tessellation_statistics(
            saved_nuclei, saved_vs, interp_depths, ax=ax
        )
        
        ax.set_title(f'Chain {chain.id}')
        ax.tick_params(direction='in', labelleft=False, labelbottom=False)
        ax.set_xlabel('')
        ax.set_ylabel('')
        
        if not ipanel % cols:
            ax.set_ylabel('Depth [km]')
            ax.tick_params(labelleft=True)
        if ipanel >= (rows-1) * cols:
            ax.set_xlabel('Vs [km/s]')
            ax.tick_params(labelbottom=True)
    
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, f"{filename}_inversion_chains.png"), dpi=300, bbox_inches='tight')
    plt.close()
    
    # 保存反演结果数据
    NS = len(sampled_thickness)
    ST = np.zeros([NS, 20],dtype=np.float32)
    SV = np.zeros([NS, 20],dtype=np.float32)
    for i in range(NS):
        ni = len(sampled_thickness[i])
        ST[i, :len(sampled_thickness[i])] = sampled_thickness[i]
        SV[i, :len(sampled_vs[i])] = sampled_vs[i]
        
    save_dict = {
        'sampled_thickness': ST,
        'sampled_vs': SV,
        'statistics_vs': statistics_vs,
        'interp_depths': interp_depths,
        'R0_periods': periods0,
        'R0_observed': d_r0,
        'R0_predicted': results['R0.dpred']
    }
    
    for i in range(0, n_curves):
        mode_data = curves_data[i]
        freqi, d_ri = mode_data[-1::-1, 0], mode_data[-1::-1, 1]/1000
        periodsi = 1.0 / freqi
        
        save_dict.update({
            f'R{i}_periods': periodsi,
            f'R{i}_observed': d_ri,
            f'R{i}_predicted': results[f'R{i}.dpred']
        })
    
    np.savez(os.path.join(output_dir, f"{filename}_inversion_results.npz"), **save_dict)
    
    print(f"文件 {filename} 的反演完成，结果已保存到 {output_dir}")

def process_all_disp_files():
    """
    处理所有频散曲线文件
    """
    # 查找所有disp文件
    # disp_files = glob.glob("./figures/6.traceSeq.figures/test/*.disp")
    # disp_files = glob.glob("script.inv/P264S.D123.Tp.N.F1.0.3.5.V100.1400.NORM.STACK.FOLDER.disp")
    # disp_files = glob.glob("script.inv/testr3.disp")
    # disp_files = glob.glob("figures/6.traceSeq.figures/test/P34x.D123.Tp.NS.F1.0.20.0.V100.1400.NORM.STACK.FOLDER.disp")
    disp_files = glob.glob("figures/6.traceSeq.figures/test/Pxxx.P192.D123.Tp.NS.F1.0.40.0.V100.1400.NORM.STACK.FOLDER.disp")
    disp_files.sort()
    print(f"找到 {len(disp_files)} 个频散曲线文件")
    disp_files = disp_files
    for i, disp_file in enumerate(disp_files):
        filename = os.path.basename(disp_file)
        print(disp_file)
        print(f"\n处理文件 ({i+1}/{len(disp_files)}): {filename}")
        
        try:
            # 读取曲线数据
            curves = read_dispersion_curves(disp_file)
            print(f"  包含 {len(curves)} 条曲线")
            
            # 创建输出目录
            output_dir = os.path.join(output_root_dir, filename.replace('.disp', ''))

            raw_h5file = disp_file.replace('.disp', '.h5')
            fv_file= raw_h5file if os.path.exists(raw_h5file) else None

            # 执行反演
            perform_inversion(curves, output_dir, filename.replace('.disp', ''),
                              fv_file=fv_file)
            
        except Exception as e:
            print(f"处理文件 {filename} 时出错: {e}")
            continue
    
    print("\n所有文件处理完成!")

if __name__ == "__main__":
    process_all_disp_files()